#!/bin/bash

###############################################
# 网络拓扑自动绘制脚本
# 名称：generate_network_topology.sh
# 版本：1.0
# 作者：资深程序员
# 日期：2025-03-23
# 描述：
#   使用nmap扫描局域网设备，并结合graphviz
#   自动绘制网络拓扑图，帮助管理员快速了解
#   网络结构和设备连接关系。
###############################################

# 获取脚本所在目录
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
LOG_FILE="$SCRIPT_DIR/network_topology.log"
OUTPUT_FILE="$SCRIPT_DIR/network_topology.png"
TEMP_FILE="$SCRIPT_DIR/nmap_result.xml"
DOT_FILE="$SCRIPT_DIR/network_topology.dot"

# 记录日志
log() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

# 检查并安装依赖
check_install() {
    local tool="$1"
    if ! command -v "$tool" &> /dev/null; then
        log "工具 $tool 未安装，正在安装..."
        sudo apt-get update && sudo apt-get install -y "$tool"
    fi
}

# 执行网络扫描
scan_network() {
    local subnet="$1"
    log "开始扫描局域网：$subnet"
    sudo nmap -sn -oX "$TEMP_FILE" "$subnet" | tee -a "$LOG_FILE"
    if [ $? -ne 0 ]; then
        log "网络扫描失败，请检查nmap是否正确安装及网络配置。"
        exit 1
    fi
    log "网络扫描完成，结果已保存到：$TEMP_FILE"
}

# 解析nmap结果并生成DOT文件
generate_dot() {
    log "开始生成DOT文件：$DOT_FILE"
    echo "graph NetworkTopology {" > "$DOT_FILE"
    echo "    layout=neato;" >> "$DOT_FILE"
    echo "    node [shape=ellipse, style=filled, color=lightblue];" >> "$DOT_FILE"

    local hosts=$(xmllint --xpath "//host[address]" "$TEMP_FILE" 2>/dev/null)
    local nodes=()
    while IFS= read -r line; do
        ip=$(echo "$line" | grep -oP 'addr="\K[0-9.]+')
        mac=$(echo "$line" | grep -oP 'addr="\K[0-9A-F:]+')
        if [[ -n "$ip" ]]; then
            if [[ -n "$mac" ]]; then
                nodes+=("\"$ip\" [label=\"$ip\\n$mac\"];")
            else
                nodes+=("\"$ip\" [label=\"$ip\"];")
            fi
        fi
    done <<< "$hosts"

    for node in "${nodes[@]}"; do
        echo "    $node" >> "$DOT_FILE"
    done

    # 简单拓扑连接
    for ((i=0; i<${#nodes[@]}-1; i++)); do
        for ((j=i+1; j<${#nodes[@]}; j++)); do
            node1=$(echo "${nodes[i]}" | grep -oP '"\K[0-9.]+')
            node2=$(echo "${nodes[j]}" | grep -oP '"\K[0-9.]+')
            echo "    \"$node1\" -- \"$node2\";" >> "$DOT_FILE"
        done
    done

    echo "}" >> "$DOT_FILE"
    log "DOT文件生成完成：$DOT_FILE"
}

# 生成拓扑图
generate_graph() {
    log "开始生成网络拓扑图：$OUTPUT_FILE"
    dot -Tpng "$DOT_FILE" -o "$OUTPUT_FILE"
    if [ $? -eq 0 ]; then
        log "网络拓扑图已成功生成：$OUTPUT_FILE"
    else
        log "生成网络拓扑图失败，请检查graphviz是否安装正确。"
    fi
}

# 主程序入口
main() {
    log "网络拓扑自动绘制脚本启动..."
    check_install "nmap"
    check_install "graphviz"
    check_install "xmllint"

    # 获取本机子网信息
    subnet=$(ip -4 route show default | grep -oP 'src \K[0-9.]+/24')
    if [[ -z "$subnet" ]]; then
        log "无法自动获取子网，请手动输入子网范围（如192.168.1.0/24）："
        read -r subnet
    fi

    scan_network "$subnet"
    generate_dot
    generate_graph
    log "网络拓扑自动绘制脚本执行完毕。"
}

main "$@"
